// prefetcher.cc
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <csignal>
#include <ctime>
#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <cerrno> 
#include <cstring> 

extern "C" {
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "prefetcher.skel.h" 
#include "prefetcher.h"
}

// ============================================================================
// 全局变量和辅助函数
// ============================================================================

std::ofstream log_file;
static volatile bool exiting = false;

// --- 新增：统计计数器 ---
static long pattern_detected_count = 0;
static long readahead_success_count = 0;
static long readahead_failure_count = 0;
static long map_failure_count = 0;

std::string get_timestamp() {
    time_t now = time(0);
    struct tm tstruct;
    char buf[80];
    tstruct = *localtime(&now);
    strftime(buf, sizeof(buf), "%Y-%m-%d %X", &tstruct);
    return buf;
}

void log_message(const std::string& msg) {
    std::string full_msg = "[" + get_timestamp() + "] " + msg;
    std::cout << full_msg << std::endl;
    if (log_file.is_open()) {
        log_file << full_msg << std::endl;
    }
}

bool find_file_offset(uint32_t pid, unsigned long addr, int& fd, off_t& file_offset) {
    // ... (函数内容保持不变) ...
    std::string maps_path = "/proc/" + std::to_string(pid) + "/maps";
    std::ifstream maps_file(maps_path);
    if (!maps_file.is_open()) {
        return false;
    }
    std::string line;
    while (std::getline(maps_file, line)) {
        std::stringstream ss(line);
        std::string address_range, perms, offset_str, dev, inode, pathname;
        ss >> address_range >> perms >> offset_str >> dev >> inode;
        std::getline(ss, pathname);
        if (!pathname.empty() && pathname[0] == ' ') {
            pathname.erase(0, pathname.find_first_not_of(" \t"));
        }
        if (pathname.empty() || pathname[0] == '[') {
            continue;
        }
        unsigned long start, end, offset;
        sscanf(address_range.c_str(), "%lx-%lx", &start, &end);
        sscanf(offset_str.c_str(), "%lx", &offset);
        if (addr >= start && addr < end) {
            fd = open(pathname.c_str(), O_RDONLY);
            if (fd < 0) {
                log_message("Error: Could not open file " + pathname);
                return false;
            }
            file_offset = (addr - start) + offset;
            return true;
        }
    }
    return false;
}

int print_libbpf_log(enum libbpf_print_level level, const char *format, va_list args) {
    char buffer[1024];
    vsnprintf(buffer, sizeof(buffer), format, args);
    log_message("libbpf: " + std::string(buffer));
    return 0;
}

static void sig_handler(int sig) {
    exiting = true;
}

// Ring buffer 事件处理函数
int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct pattern_info *e = (struct pattern_info*)data;
    
    // --- 更新计数器 ---
    pattern_detected_count++;

    std::stringstream log_msg;
    log_msg << "Pattern detected for PID: " << e->pid << ", TID: " << e->tid
            << ", Addr: 0x" << std::hex << e->start_addr
            << ", Stride: " << std::dec << e->stride
            << ", Prefetch Size: " << e->size;
    log_message(log_msg.str());

    int fd = -1;
    off_t file_offset = 0;
    if (find_file_offset(e->pid, e->start_addr, fd, file_offset)) {
        log_message("  -> Mapped to file. FD: " + std::to_string(fd) + ", Offset: " + std::to_string(file_offset));
        
        long ret = syscall(SYS_readahead, fd, file_offset, e->size);

        if (ret == 0) {
            log_message("  -> readahead(...) syscall succeeded.");
            readahead_success_count++; // 更新成功计数
        } else {
            log_message("  -> readahead() syscall failed. Return code: " + std::to_string(ret) + ", errno: " + std::to_string(errno));
            readahead_failure_count++; // 更新失败计数
        }
        close(fd);
    } else {
        log_message("  -> Failed to map address to file offset.");
        map_failure_count++; // 更新地址转换失败计数
    }
    return 0;
}


// ============================================================================
// Main 函数
// ============================================================================
int main(int argc, char **argv) {
    // ... (main函数的前半部分保持不变) ...
    if (argc > 1 && (std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help")) {
        std::cout << "Usage: " << argv[0] << " [target_pid]" << std::endl;
        return 0;
    }
    log_file.open("prefetcher.log", std::ios::out | std::ios::trunc);
    log_message("Prefetcher starting...");
    struct prefetcher_bpf *skel;
    int err;
    libbpf_set_print(print_libbpf_log);
    skel = prefetcher_bpf__open();
    if (!skel) {
        log_message("Error: Failed to open BPF skeleton");
        return 1;
    }
    err = prefetcher_bpf__load(skel);
    if (err) {
        log_message("Error: Failed to load BPF skeleton. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }
    struct config cfg = {};
    if (argc > 1) {
        cfg.target_pid = std::stoi(argv[1]);
        log_message("Targeting PID: " + std::string(argv[1]));
    } else {
        cfg.target_pid = 0;
        log_message("Monitoring all processes.");
    }
    int map_fd = bpf_map__fd(skel->maps.config_map);
    const int key = 0;
    err = bpf_map_update_elem(map_fd, &key, &cfg, BPF_ANY);
    if (err) {
        log_message("Error: Failed to update config map. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }
    err = prefetcher_bpf__attach(skel);
    if (err) {
        log_message("Error: Failed to attach BPF skeleton. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }
    log_message("eBPF program loaded and attached successfully.");
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);
    struct ring_buffer *rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event, NULL, NULL);
    if (!rb) {
        log_message("Error: Failed to create ring buffer");
        err = -1;
        goto cleanup;
    }
    log_message("Waiting for events... Press Ctrl+C to exit.");
    while (!exiting) {
        err = ring_buffer__poll(rb, 100 /* timeout, ms */);
        if (err == -EINTR) {
            err = 0;
            break;
        }
        if (err < 0) {
            log_message("Error polling ring buffer: " + std::to_string(err));
            break;
        }
    }

cleanup:
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);
    
    // --- 新增：打印最终的统计信息 ---
    log_message("Prefetcher stopped. Final statistics:");
    std::stringstream final_stats;
    final_stats << "  - Total patterns detected (page faults handled): " << pattern_detected_count << std::endl;
    final_stats << "  - Successful readahead calls: " << readahead_success_count << std::endl;
    final_stats << "  - Failed readahead calls: " << readahead_failure_count << std::endl;
    final_stats << "  - Address mapping failures: " << map_failure_count;
    log_message(final_stats.str());

    if (log_file.is_open()) {
        log_file.close();
    }
    return err < 0 ? -err : 0;
}